import torch
from transformers import BertModel, BertTokenizer
# from torch.utils.tensorboard import SummaryWriter
# writer = SummaryWriter(log_dir='logs')

model_name = '../bert-base-chinese'
model_path = '../bert-base-chinese'

tokenizer = BertTokenizer.from_pretrained(model_name)
bert_model = BertModel.from_pretrained(model_path)



# 输入文本
text = "你无敌了，我没事"
sent_id = tokenizer.encode(text,
                           add_special_tokens=True,
                           # 添加 [CLS] 和 [SEP] 标记
                           truncation=True,
                           # 指定序列的最大长度
                           max_length=10
                           )
# 打印整数序列
print("整数序列: {}".format(sent_id))
# 将整数转换回文本
print("标记化文本:",tokenizer.convert_ids_to_tokens(sent_id))

# print(bert_model)   # 打印模型
decoded = tokenizer.decode(sent_id)
print("解码字符串: {}".format(decoded))

att_mask = [int(tok > 0) for tok in sent_id]
print("注意力掩码:",att_mask)

# 将列表转换为张量
sent_id = torch.tensor(sent_id)
att_mask = torch.tensor(att_mask)
# 将张量调整为（批量大小，文本长度）的形式
# unsqueeze增加一个维度
# squeeze删掉一个维度
sent_id = sent_id.unsqueeze(0)
att_mask = att_mask.unsqueeze(0)
print("sent_id:",sent_id)
print("att_mask:",att_mask)

outputs = bert_model(sent_id, attention_mask=att_mask)

# print("outputs:",outputs)
# 解包 BERT 模型的输出
# 每个时间步的隐藏状态
all_hidden_states = outputs[0]
# 第一个时间步的隐藏状态（[CLS] 标记）
cls_hidden_state = outputs[1]
print("最后一个隐藏状态的形状:",all_hidden_states.shape)
print("CLS 隐藏状态的形状:",cls_hidden_state.shape)